#include <stdio.h>
#include <unistd.h>
#include <stdint.h>

#include "cuda_runtime.h"
#include "nccl.h"

#define CUDACHECK(cmd) do {                         \
    cudaError_t e = cmd;                              \
    if( e != cudaSuccess ) {                          \
        printf("Failed: Cuda error %s:%d '%s'\n",             \
                __FILE__,__LINE__,cudaGetErrorString(e));   \
        exit(EXIT_FAILURE);                             \
    }                                                 \
} while(0)

#define NCCLCHECK(cmd) do {                         \
    ncclResult_t r = cmd;                             \
    if (r!= ncclSuccess) {                            \
        printf("Failed, NCCL error %s:%d '%s'\n",             \
                __FILE__,__LINE__,ncclGetErrorString(r));   \
        exit(EXIT_FAILURE);                             \
    }                                                 \
} while(0)
#define MIN_ARGC                        (1)
#define MAX_SUPPORT_DEV_NUM_COUNT       (8)
char sCmdGetGPUNumber[] = "nvidia-smi topo -m";
int aSupportDevNum[MAX_SUPPORT_DEV_NUM_COUNT] = {1, 2, 4, 8};
int main(int argc, char* argv[])
{
    //each process is using x GPUs
    int nDev = 1;
    int nRanks = nDev;
    ncclDataType_t datatype = ncclInt32;

    int nCount = 4;
    int nTotalCount = nDev * nCount;
    if(1 == argc)
    {
        printf("[%s +%d %s] argc=%d !=%d so exit(0)\n", __FILE__, __LINE__, __func__, argc, MIN_ARGC);
        printf("[%s +%d %s] example ./nccl_send_recv.elf num is the number can used \n", __FILE__, __LINE__, __func__);
        printf("./nccl_send_recv.elf 1 \n");
        printf("./nccl_send_recv.elf 2 \n");
        printf("./nccl_send_recv.elf 4 \n");
        system(sCmdGetGPUNumber);
        exit(0);
    }
    else
    {
      int i = 0;
      nDev = atoi(argv[1]);
      printf("[%s +%d %s] input nDev=%d\n", __FILE__, __LINE__, __func__, nDev);
      printf("[%s +%d %s] MAX_SUPPORT_DEV_NUM_COUNT=%d\n", __FILE__, __LINE__, __func__, MAX_SUPPORT_DEV_NUM_COUNT);
      for(i = 0; i < MAX_SUPPORT_DEV_NUM_COUNT; i++)
      {
        if(nDev == aSupportDevNum[i])
        {
          break;
        }
      }
      if(MAX_SUPPORT_DEV_NUM_COUNT == i)
      {
          printf("[%s +%d %s] nDev=%d is not in aSupportDevNum[] so exit(0)\n", __FILE__, __LINE__, __func__, nDev);
          printf("[%s +%d %s] example ./nccl_send_recv.elf num is the number can used \n", __FILE__, __LINE__, __func__);
          printf("./nccl_send_recv.elf 1 \n");
          printf("./nccl_send_recv.elf 2 \n");
          printf("./nccl_send_recv.elf 4 \n");
          printf("./nccl_send_recv.elf 8 \n");
          system(sCmdGetGPUNumber);
          exit(0);
      }
    }
    nRanks = nDev;
    nTotalCount = nDev * nCount;

    int32_t** devSendBuff = (int32_t**)malloc(nDev * sizeof(int32_t*));
    int32_t** devRecvBuff = (int32_t**)malloc(nDev * sizeof(int32_t*));
    cudaStream_t* psStream = (cudaStream_t*)malloc(sizeof(cudaStream_t)*nDev);

    int32_t** hostSendBuff = (int32_t**)malloc(nDev * sizeof(int32_t*));
    int32_t** hostRecvBuff = (int32_t**)malloc(nDev * sizeof(int32_t*));

  printf("=====================start==============\n");
  printf("[%s +%d %s] ncclSend+ncclRecv API test example\n", __FILE__, __LINE__, __func__);

    //picking GPUs based on localRank
    for (int i = 0; i < nDev; ++i) {
        CUDACHECK(cudaSetDevice(i));
        CUDACHECK(cudaMalloc(devSendBuff + i, nTotalCount * sizeof(int32_t)));
        CUDACHECK(cudaMalloc(devRecvBuff + i, nTotalCount * sizeof(int32_t)));
        hostSendBuff[i] = (int32_t*)malloc(nCount * sizeof(int32_t));
        hostRecvBuff[i] = (int32_t*)malloc(nCount * sizeof(int32_t));

        CUDACHECK(cudaStreamCreate(psStream+i));

        #if 0
        CUDACHECK(cudaMemset(devSendBuff[i], 1, nTotalCount * sizeof(int32_t)));
        CUDACHECK(cudaMemset(devRecvBuff[i], 0, nTotalCount * sizeof(int32_t)));
        #else
        #endif

    }
    for (int i = 0; i < nDev; ++i) {
        for (int j = 0; j < nCount; ++j)
        {
            hostSendBuff[i][j] = (i + 1) * (j + 1) * 10;
            hostRecvBuff[i][j] = 1;
            printf("[%s +%d %s] hostSendBuff[i=%d][j=%d]=%d\n", __FILE__, __LINE__, __func__, i, j, hostSendBuff[i][j]);
            printf("[%s +%d %s] hostRecvBuff[i=%d][j=%d]=%d\n", __FILE__, __LINE__, __func__, i, j, hostRecvBuff[i][j]);
        }
        CUDACHECK(cudaMemcpy(devSendBuff[i], hostSendBuff[i], nCount * sizeof(int32_t), cudaMemcpyHostToDevice));
        CUDACHECK(cudaMemcpy(devRecvBuff[i], hostRecvBuff[i], nCount * sizeof(int32_t), cudaMemcpyHostToDevice));
    }

    ncclUniqueId sID;
    ncclComm_t asComms[nDev];
    //generating NCCL unique sID at one process and broadcasting it to all
    ncclGetUniqueId(&sID);

    //initializing NCCL, group API is required around ncclCommInitRank as it is
    //called across multiple GPUs in each thread/process
    NCCLCHECK(ncclGroupStart());
    for (int i=0; i<nDev; i++) {
        CUDACHECK(cudaSetDevice(i));
        NCCLCHECK(ncclCommInitRank(asComms+i, nRanks, sID, i));
    }
    NCCLCHECK(ncclGroupEnd());

    printf("[%s +%d %s] after call ncclCommInitRank\n", __FILE__, __LINE__, __func__);

    //calling NCCL communication API. Group API is required when using
    //multiple devices per thread/process
    NCCLCHECK(ncclGroupStart());
    for (int i=0; i<nDev; i++) {
        printf("[%s +%d %s] ncclSend parameter4 peer(i)=%d\n", __FILE__, __LINE__, __func__, i);
        NCCLCHECK(ncclSend((const void*)(devSendBuff[i]), nCount, datatype, i, asComms[i], psStream[i]));
        printf("[%s +%d %s] ncclRecv parameter4 peer(i)=%d\n", __FILE__, __LINE__, __func__, i);
        NCCLCHECK(ncclRecv((void*)(devRecvBuff[i]), nCount, datatype, i, asComms[i], psStream[i]));
    }
    NCCLCHECK(ncclGroupEnd());

    //synchronizing on CUDA stream to complete NCCL communication
    for (int i=0; i<nDev; i++)
    {
        CUDACHECK(cudaSetDevice(i));
        CUDACHECK(cudaStreamSynchronize(psStream[i]));
    }
    for (int i = 0; i < nDev; ++i) {
        CUDACHECK(cudaSetDevice(i));
        CUDACHECK(cudaMemcpy(hostSendBuff[i], devSendBuff[i], nCount * sizeof(int32_t), cudaMemcpyDeviceToHost));
        CUDACHECK(cudaMemcpy(hostRecvBuff[i], devRecvBuff[i], nCount * sizeof(int32_t), cudaMemcpyDeviceToHost));
    }
    printf("[%s +%d %s] after ncclSend+ncclRecv+cudaStreamSynchronize+cudaMemcpy\n", __FILE__, __LINE__, __func__);
    for (int i = 0; i < nDev; ++i) {
        for (int j = 0; j < nCount; ++j)
        {
            printf("[%s +%d %s] hostSendBuff[i=%d][j=%d]=%d\n", __FILE__, __LINE__, __func__, i, j, hostSendBuff[i][j]);
            printf("[%s +%d %s] hostRecvBuff[i=%d][j=%d]=%d\n", __FILE__, __LINE__, __func__, i, j, hostRecvBuff[i][j]);
        }
    }

    //freeing device memory
    for (int i=0; i<nDev; i++) {
        CUDACHECK(cudaFree(devSendBuff[i]));
        CUDACHECK(cudaFree(devRecvBuff[i]));
    }
    if(NULL != devSendBuff)
    {
        free(devSendBuff);
        devSendBuff = NULL;
    }
    if(NULL != devRecvBuff)
    {
        free(devRecvBuff);
        devRecvBuff = NULL;
    }   
    
    for (int i=0; i<nDev; i++) {
        free(hostSendBuff[i]);
        free(hostRecvBuff[i]);
    }
    if(NULL != hostSendBuff)
    {
        free(hostSendBuff);
        hostSendBuff = NULL;
    }
    if(NULL != hostRecvBuff)
    {
        free(hostRecvBuff);
        hostRecvBuff = NULL;
    }
    
    if(NULL != psStream)
    {
        free(psStream);
        psStream = NULL;
    }
    //finalizing NCCL
    for (int i=0; i<nDev; i++) {
        ncclCommDestroy(asComms[i]);
    }
    printf("=====================end==============\n");

    return 0;
}
